package com.dddframework.web.utils;

import com.dddframework.core.utils.JsonKit;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketHttpHeaders;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.client.WebSocketClient;
import org.springframework.web.socket.client.WebSocketConnectionManager;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.net.URI;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Consumer;

/**
 * WebSocket工具
 */
@Slf4j(topic = "### BASE-WEB : BaseWebSocketClient ###")
public class BaseWebSocketClient {
    private WebSocketClient webSocketClient;

    public BaseWebSocketClient(WebSocketClient webSocketClient) {
        this.webSocketClient = webSocketClient;
    }

    private static final Map<String, WebSocketSession> SESSIONS = new ConcurrentHashMap<>();

    public void connectServer(String biz, String url, Consumer<String> listener) {
        TextWebSocketHandler handler = new TextWebSocketHandler() {
            // 处理连接后保存session
            @Override
            public void afterConnectionEstablished(WebSocketSession session) {
                log.info("WebSocket is established.");
                SESSIONS.put(biz, session);
            }

            @Override
            protected void handleTextMessage(WebSocketSession session, TextMessage message) {
                log.info("Received WebSocket message: {}", message.getPayload());
                try {
                    listener.accept(message.getPayload());
                } catch (Exception e) {
                    log.error("Handle WebSocket message failed!", e);
                }
            }

            @Override
            public void handleTransportError(WebSocketSession session, Throwable exception) {
                log.error("WebSocket transport error: {}", exception.getMessage(), exception);
            }
        };
        try {
            webSocketClient.doHandshake(handler, new WebSocketHttpHeaders(), new URI(url)).get();
        } catch (Exception e) {
            log.error("Failed to connect to WebSocket server at {}: {}", url, e.getMessage(), e);
        }
        new WebSocketConnectionManager(webSocketClient, handler, url).start();
    }

    public void connectServer(String url, Consumer<String> listener) {
        connectServer("DEFAULT", url, listener);
    }

    @SneakyThrows
    public void disconnect(String biz) {
        WebSocketSession session = SESSIONS.get(biz);
        if (session != null && session.isOpen()) {
            SESSIONS.remove(biz);
            log.info("WebSocket session is close.");
        }
    }

    @SneakyThrows
    public static void sendMessage(String biz, Object message) {
        WebSocketSession webSocketSession = SESSIONS.get(biz);
        if (webSocketSession != null && webSocketSession.isOpen()) {
            webSocketSession.sendMessage(new TextMessage(JsonKit.toJson(message)));
        } else {
            log.error("WebSocket is not connected.");
        }
    }

    public static void sendMessage(Object message) {
        sendMessage("DEFAULT", message);
    }
}